import smtplib
import logging
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.mime.application import MIMEApplication
import os
from datetime import datetime
from app.models.database import db, Scan, Schedule, User
from app.services.report_service import ReportService
from app import app

logger = logging.getLogger(__name__)

class NotificationService:
    """通知服务，负责扫描结果通知和报告发送"""
    
    @staticmethod
    def send_scan_completion_notification(scan_id):
        """
        发送扫描完成通知
        
        Args:
            scan_id (str): 扫描任务ID
        """
        try:
            # 获取扫描信息
            scan = Scan.query.get(scan_id)
            if not scan:
                logger.error(f"未找到扫描任务: {scan_id}")
                return False
            
            # 如果是计划任务，获取关联的调度信息和用户信息
            if scan.schedule_id:
                schedule = Schedule.query.get(scan.schedule_id)
                if schedule:
                    user = User.query.get(schedule.user_id)
                    if user and (user.email or schedule.notification_email):
                        # 获取漏洞统计信息
                        from app.services.scan_service import get_scan_summary
                        summary = get_scan_summary(scan_id)
                        
                        # 生成报告
                        report_format = schedule.report_format or 'pdf'
                        report_content, extension, _ = ReportService.generate_report(scan_id, report_format)
                        
                        # 发送邮件
                        recipient = schedule.notification_email or user.email
                        NotificationService._send_scan_report_email(
                            recipient=recipient,
                            scan=scan,
                            summary=summary,
                            report_content=report_content,
                            report_format=report_format
                        )
                        logger.info(f"已发送扫描报告邮件到: {recipient}")
                        return True
            
            # 对于非计划任务，检查用户关联
            if scan.user_id:
                user = User.query.get(scan.user_id)
                if user and user.email:
                    # 发送基本通知
                    NotificationService._send_scan_completion_email(
                        recipient=user.email,
                        scan=scan
                    )
                    logger.info(f"已发送扫描完成通知到: {user.email}")
                    return True
                    
        except Exception as e:
            logger.error(f"发送扫描完成通知失败: {str(e)}")
        
        return False
    
    @staticmethod
    def _send_scan_report_email(recipient, scan, summary, report_content, report_format):
        """
        发送包含报告附件的扫描结果邮件
        """
        # 构建邮件内容
        subject = f"【安全扫描】{scan.target} 扫描完成报告"
        body = f"""尊敬的用户：

您的安全扫描任务已完成，详情如下：

扫描目标: {scan.target}
扫描时间: {scan.created_at.strftime('%Y-%m-%d %H:%M:%S')}
完成时间: {scan.finished_at.strftime('%Y-%m-%d %H:%M:%S') if scan.finished_at else '未知'}

漏洞统计：
- 总漏洞数: {summary.get('total', 0)}
- 高危漏洞: {summary.get('high', 0)}
- 中危漏洞: {summary.get('medium', 0)}
- 低危漏洞: {summary.get('low', 0)}

完整的扫描报告已作为附件提供，请下载查看详情。

如有问题，请联系系统管理员。

此致，
安全扫描系统
        """
        
        # 创建邮件
        msg = MIMEMultipart()
        msg['From'] = app.config['MAIL_USERNAME']
        msg['To'] = recipient
        msg['Subject'] = subject
        
        # 添加正文
        msg.attach(MIMEText(body, 'plain', 'utf-8'))
        
        # 添加附件
        filename = f"scan_report_{scan.id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.{report_format}"
        
        if isinstance(report_content, bytes):
            attachment = MIMEApplication(report_content)
        else:
            attachment = MIMEText(report_content, 'plain', 'utf-8')
        
        attachment.add_header('Content-Disposition', 'attachment', filename=filename)
        msg.attach(attachment)
        
        # 发送邮件
        NotificationService._send_email(msg)
    
    @staticmethod
    def _send_scan_completion_email(recipient, scan):
        """
        发送简单的扫描完成通知邮件
        """
        subject = f"【安全扫描】{scan.target} 扫描任务已完成"
        body = f"""尊敬的用户：

您的安全扫描任务已完成，详情如下：

扫描目标: {scan.target}
扫描ID: {scan.id}
扫描状态: {scan.status}
扫描时间: {scan.created_at.strftime('%Y-%m-%d %H:%M:%S')}

请登录系统查看详细的扫描结果。

此致，
安全扫描系统
        """
        
        # 创建邮件
        msg = MIMEMultipart()
        msg['From'] = app.config['MAIL_USERNAME']
        msg['To'] = recipient
        msg['Subject'] = subject
        
        # 添加正文
        msg.attach(MIMEText(body, 'plain', 'utf-8'))
        
        # 发送邮件
        NotificationService._send_email(msg)
    
    @staticmethod
    def _send_email(msg):
        """
        发送邮件的通用方法
        """
        try:
            # 连接SMTP服务器
            if app.config['MAIL_USE_SSL']:
                server = smtplib.SMTP_SSL(app.config['MAIL_SERVER'], app.config['MAIL_PORT'])
            else:
                server = smtplib.SMTP(app.config['MAIL_SERVER'], app.config['MAIL_PORT'])
                if app.config['MAIL_USE_TLS']:
                    server.starttls()
            
            # 登录
            server.login(app.config['MAIL_USERNAME'], app.config['MAIL_PASSWORD'])
            
            # 发送邮件
            text = msg.as_string()
            server.sendmail(app.config['MAIL_USERNAME'], msg['To'], text)
            
            # 关闭连接
            server.quit()
            
            logger.info(f"邮件发送成功: 收件人={msg['To']}, 主题={msg['Subject']}")
            
        except Exception as e:
            logger.error(f"邮件发送失败: {str(e)}")
            raise
    
    @staticmethod
    def send_high_vulnerability_alert(scan_id, high_vulnerabilities):
        """
        当发现高危漏洞时发送警报
        
        Args:
            scan_id (str): 扫描任务ID
            high_vulnerabilities (list): 高危漏洞列表
        """
        try:
            # 获取扫描信息
            scan = Scan.query.get(scan_id)
            if not scan:
                return False
            
            # 获取相关用户
            users = []
            
            # 检查是否关联了计划任务
            if scan.schedule_id:
                schedule = Schedule.query.get(scan.schedule_id)
                if schedule:
                    user = User.query.get(schedule.user_id)
                    if user and user.email:
                        users.append(user.email)
            
            # 检查是否关联了用户
            if scan.user_id:
                user = User.query.get(scan.user_id)
                if user and user.email and user.email not in users:
                    users.append(user.email)
            
            # 添加管理员邮箱（如果配置了）
            if 'ADMIN_EMAIL' in app.config and app.config['ADMIN_EMAIL'] and app.config['ADMIN_EMAIL'] not in users:
                users.append(app.config['ADMIN_EMAIL'])
            
            # 发送警报邮件
            for recipient in users:
                NotificationService._send_high_vuln_alert_email(
                    recipient=recipient,
                    scan=scan,
                    high_vulnerabilities=high_vulnerabilities
                )
            
            return True
            
        except Exception as e:
            logger.error(f"发送高危漏洞警报失败: {str(e)}")
        
        return False
    
    @staticmethod
    def _send_high_vuln_alert_email(recipient, scan, high_vulnerabilities):
        """
        发送高危漏洞警报邮件
        """
        subject = f"【高危漏洞警报】{scan.target} 发现高危安全漏洞"
        
        # 构建高危漏洞列表
        vuln_list = ""
        for vuln in high_vulnerabilities:
            vuln_list += f"\n- 类型: {vuln.get('vuln_type', '未知')}"
            vuln_list += f"\n  路径: {vuln.get('path', '未知')}"
            vuln_list += f"\n  风险评分: {vuln.get('risk_score', '未知')}"
            vuln_list += f"\n  描述: {vuln.get('description', '无描述')}\n"
        
        body = f"""⚠️ 高危漏洞警报 ⚠️

尊敬的用户：

在对目标系统 {scan.target} 的安全扫描中，发现了以下高危漏洞，需要立即处理：

{vuln_list}

建议立即查看详细报告并进行修复。

此致，
安全扫描系统
        """
        
        # 创建并发送邮件
        msg = MIMEMultipart()
        msg['From'] = app.config['MAIL_USERNAME']
        msg['To'] = recipient
        msg['Subject'] = subject
        msg.attach(MIMEText(body, 'plain', 'utf-8'))
        
        NotificationService._send_email(msg)